# coding=utf-8
# Copyright (C) idata team - All Rights Reserved
#
# @Version:   3.10.9
# @Software:  PyCharm
# @FileName:  metrics.py
# @CTime:     2023/6/2 15:24   
# @Author:    yhy
# @Email:     yhy@cyber.com
# @UTime:     2023/6/2 15:24
#
# @Description:
#     
#     xxx
#
import logging
from typing import NewType, Any, Optional
import torch
from torchmetrics import Metric

logger = logging.getLogger(__name__)


class MyAccuracy(Metric):
    def __init__(self, dist_sync_on_step=False):
        # call `self.add_state`for every internal state that is needed for the metrics computations
        # dist_reduce_fx indicates the function that should be used to reduce state from multiple processes
        super().__init__(dist_sync_on_step=dist_sync_on_step)

        self.add_state("correct", default=torch.tensor(0), dist_reduce_fx="sum")
        self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")

    def update(self, preds: torch.Tensor, target: torch.Tensor):
        # update metric states
        preds, target = self._input_format(preds, target)
        assert preds.shape == target.shape

        self.correct += torch.sum(preds == target)
        self.total += target.numel()

    def compute(self):
        # compute final result
        return self.correct.float() / self.total



if __name__ == '__main__':
    main()
